import logging

from datasets.specific_pair_match_datasets import *
from datasets.specific_temp_match_datasets import *
from tools.json_helper import json_helper


logger = logging.getLogger("match")


class ComprehensiveDataset:
    cls_dict = {
        "key/1" : None, 
        "key/2" : None, 
        "key/3" : None
    }
    description = "Comprehensive"

    def __init__(self, config=None):
        '''
        config.list_file is a dict of different dataset list files.
        keys correspond to name in cls_dict, 
        if there are repeated dataset, append suffix like '/+suffix'
        '''
        list_file = config.list_file
        cls_call = json_helper(json_dir=list_file).get_dict()

        self.len_list = []
        self.dset_list = []
        for k, v in cls_call.items():
            k = k.split("/")[0]
            cls_ = self.cls_dict[k]
            config.list_file = v
            self.dset_list.append(cls_(config))

        config.list_file = list_file

        logger.info(f"{self.description} Dataset created with {self.__len__()} pairs")


    def __len__(self):
        length = 0
        self.len_list.clear()
        for d in self.dset_list:
            length += len(d)
            self.len_list.append(length)

        return length

    
    def __getitem__(self, index):
        dset_ids = 0
        for dset_ids, l in enumerate(self.len_list):
            if l > index:
                break

        len_base = 0 if dset_ids == 0 else self.len_list[dset_ids-1]
        index -= len_base

        ret = self.dset_list[dset_ids].__getitem__(index)

        return ret


class ComprehensivePairMDataset(ComprehensiveDataset):
    cls_dict = {
        "Sentinel12" : Sentinel12PairMDataset, 
        "OS" : OSPairMDataset, 
        "QXSLAB" : QXSLABPairMDataset
    }
    description = "Comprehensive Pair Match"


class ComprehensiveTempMDataset(ComprehensiveDataset):
    cls_dict = {
        "Sentinel12" : Sentinel12PairDataset, 
        "OS" : OSPairDataset, 
        "QXSLAB" : QXSLABPairDataset
    }
    description = "Comprehensive Template Match"


# test ComprehensivePairMDataset
if __name__ == "__main__1":
    import cv2 as cv
    from dotmap import DotMap
    import numpy as np

    dset_kwt = "ComprehensivePairMDataset"
    list_file = "E:\workspace\SOMatch\json\comprehensive\pair_match_train.json"
    dset_obj = ComprehensivePairMDataset

    ids = [100, 2000, 6000, 8000]

    class config:

        def __init__(self):
            self.single_domain = ""
            self.cache_dir = None
            self.cache_size = 0
            self.base_dir = None
            self.list_file = list_file
            self.augment = {
                "fliplr": 0,
                "flipud": 0,
                "scale": 0,
                "scale_px": (1.0, 1.0),
                "translate": 0,
                "translate_perc": (0.0, 0.0),
                "rotate": 0,
                "rotate_angle": (-5, 5)
            }
            self.clipper = DotMap({
                    "name": "center same",
                    "template_size": 128
            })

    cfg = config()

    ppd = dset_obj(cfg)

    for i in ids:

        img0, img1 = ppd[i]
        img0 = img0.numpy()
        img1 = img1.numpy()
        img0 = np.transpose(img0, [1, 2, 0])
        img1 = np.transpose(img1, [1, 2, 0])

        cv.imwrite(f"image/dataset_test/{dset_kwt}_{i}_pariM_sar.png", img0*255)
        cv.imwrite(f"image/dataset_test/{dset_kwt}_{i}_pariM_opt.png", img1*255)

    # print(ppd.len_list)


# test ComprehensiveTempMDataset
if __name__ == "__main__":
    import cv2 as cv
    from dotmap import DotMap
    import numpy as np

    dset_kwt = "ComprehensiveTempMDataset"
    list_file = "E:\workspace\SOMatch\json\comprehensive\pair_match_train.json"
    dset_obj = ComprehensiveTempMDataset

    ids = [100, 2000, 6000, 8000]

    class config:

        def __init__(self):
            self.single_domain = ""
            self.cache_dir = None
            self.cache_size = 0
            self.base_dir = None
            self.list_file = list_file
            self.augment = {
                "fliplr": 0,
                "flipud": 0,
                "scale": 0,
                "scale_px": (1.0, 1.0),
                "translate": 0,
                "translate_perc": (0.0, 0.0),
                "rotate": 0,
                "rotate_angle": (-5, 5)
            }
            self.clipper = DotMap({
                    "name": "shift template",
                    "template_size": 144,
                    "shift_x": 15,
                    "shift_y": 15,
                    "search_rad": 32,
                    "search_domain": "A"
            })

    cfg = config()

    ppd = dset_obj(cfg)

    for i in ids:

        (img0, img1), (gt_map, _) = ppd[i]
        img0 = img0.numpy()
        img1 = img1.numpy()
        gt_map = gt_map.numpy()
        img0 = np.transpose(img0, [1, 2, 0])
        img1 = np.transpose(img1, [1, 2, 0])

        cv.imwrite(f"image/dataset_test/{dset_kwt}_{i}_sar.png", img0*255)
        cv.imwrite(f"image/dataset_test/{dset_kwt}_{i}_opt.png", img1*255)
        cv.imwrite(f"image/dataset_test/{dset_kwt}_{i}_map.png", gt_map*255)

    # print(ppd.len_list)
